package cn.demoncat.util.sec;

import cn.demoncat.util.io.CloseUtil;

import javax.crypto.Cipher;
import java.io.FileInputStream;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;

/**
 * 基于OpenSSL密钥库和数字证书的加密解密和签名验证
 * <br/>
 * 服务端掌管密钥库和私钥；客户端掌管分发的数字证书和公钥
 */
public class CertificateUtil extends SecurityUtil {

	/**
	 * 私钥加密
	 * 
	 * @param data 原数据
	 * @param alias 证书别名
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	public static byte[] encryptByPri(byte[] data, String alias, String keyStorePath, String keyStorePwd, boolean isBoth) {
		// 获取私钥
		PrivateKey privateKey = getPrivateKeyByKeyStore(alias, keyStorePath, keyStorePwd, isBoth);
		// 对数据加密
		try {
			Cipher cipher = Cipher.getInstance(privateKey.getAlgorithm(),PROVIDER);
			cipher.init(Cipher.ENCRYPT_MODE, privateKey);
			return cipher.doFinal(data);
		} catch (Exception e) {
			throw new RuntimeException("私钥加密失败",e);
		}
	}

	/**
	 * 私钥加密
	 * 
	 * @param data 原数据
	 * @param alias 证书别名
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	public static String encryptByPri(String data, String alias, String keyStorePath, String keyStorePwd, boolean isBoth) {
		try {
			return Base64Util.encode(encryptByPri(data.getBytes(StandardCharsets.UTF_8), alias, keyStorePath, keyStorePwd, isBoth));
		} catch (Exception e) {
			throw new RuntimeException("私钥加密失败",e);
		}
	}

	/**
	 * 私钥解密
	 * 
	 * @param data 加密数据
	 * @param alias 证书别名
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	public static byte[] decryptByPri(byte[] data, String alias, String keyStorePath, String keyStorePwd, boolean isBoth) {
		// 取得私钥
		PrivateKey privateKey = getPrivateKeyByKeyStore(alias, keyStorePath, keyStorePwd, isBoth);
		// 对数据解密
		try {
			Cipher cipher = Cipher.getInstance(privateKey.getAlgorithm(),PROVIDER);
			cipher.init(Cipher.DECRYPT_MODE, privateKey);
			return cipher.doFinal(data);
		} catch (Exception e) {
			throw new RuntimeException("私钥解密失败",e);
		}
	}

	/**
	 * 私钥解密
	 * 
	 * @param data 加密数据
	 * @param alias 证书别名
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	public static String decryptByPri(String data, String alias, String keyStorePath, String keyStorePwd, boolean isBoth) {
		try {
			return new String(decryptByPri(Base64Util.decode(data), alias, keyStorePath, keyStorePwd, isBoth), StandardCharsets.UTF_8);
		} catch (Exception e) {
			throw new RuntimeException("私钥解密失败",e);
		}
	}

	/**
	 * 公钥加密
	 * 
	 * @param data 原数据
	 * @param certificatePath 证书路径
	 */
	public static byte[] encryptByPub(byte[] data, String certificatePath) {
		// 取得公钥
		PublicKey publicKey = getPublicKeyByCertificate(certificatePath);
		// 对数据加密
		try {
			Cipher cipher = Cipher.getInstance(publicKey.getAlgorithm(),PROVIDER);
			cipher.init(Cipher.ENCRYPT_MODE, publicKey);
			return cipher.doFinal(data);
		} catch (Exception e) {
			throw new RuntimeException("公钥加密失败",e);
		}
	}

	/**
	 * 公钥加密
	 * 
	 * @param data 原数据
	 * @param certificatePath 证书路径
	 */
	public static String encryptByPub(String data, String certificatePath) {
		try {
			return Base64Util.encode(encryptByPub(data.getBytes(StandardCharsets.UTF_8), certificatePath));
		} catch (Exception e) {
			throw new RuntimeException("公钥加密失败",e);
		}
	}

	/**
	 * 公钥解密
	 * 
	 * @param data 加密数据
	 * @param certificatePath 证书路径
	 */
	public static byte[] decryptByPub(byte[] data, String certificatePath) {
		// 取得公钥
		PublicKey publicKey = getPublicKeyByCertificate(certificatePath);
		// 对数据解密
		try {
			Cipher cipher = Cipher.getInstance(publicKey.getAlgorithm(),PROVIDER);
			cipher.init(Cipher.DECRYPT_MODE, publicKey);
			return cipher.doFinal(data);
		} catch (Exception e) {
			throw new RuntimeException("公钥解密失败",e);
		}
	}

	/**
	 * 公钥解密
	 * 
	 * @param data 加密数据
	 * @param certificatePath 证书路径
	 */
	public static String decryptByPub(String data, String certificatePath) {
		try {
			return new String(decryptByPub(Base64Util.decode(data), certificatePath), StandardCharsets.UTF_8);
		} catch (Exception e) {
			throw new RuntimeException("公钥解密失败",e);
		}
	}

	/**
	 * 私钥签名
	 * 
	 * @param data 原数据
	 * @param alias 证书别名
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	public static byte[] sign(byte[] data, String alias, String keyStorePath, String keyStorePwd, boolean isBoth) {
		// 获得证书（通过密钥库）
		X509Certificate x509 = (X509Certificate) (getCertificate(alias, keyStorePath, keyStorePwd, isBoth));
		try {
			// 构建签名，由证书指定签名算法
			Signature signature = Signature.getInstance(x509.getSigAlgName(),PROVIDER);
			// 获得私钥
			PrivateKey privateKey = getPrivateKeyByKeyStore(alias, keyStorePath, keyStorePwd, isBoth);
			// 私钥签名
			signature.initSign(privateKey);
			signature.update(data);
			return signature.sign();
		} catch (Exception e) {
			throw new RuntimeException("私钥签名失败",e);
		}
	}

	/**
	 * 私钥签名
	 * 
	 * @param data 原数据
	 * @param alias 证书别名
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	public static String sign(String data, String alias, String keyStorePath, String keyStorePwd, boolean isBoth) {
		try {
			return Base64Util.encode(sign(data.getBytes(StandardCharsets.UTF_8), alias, keyStorePath, keyStorePwd, isBoth));
		} catch (Exception e) {
			throw new RuntimeException("私钥签名失败",e);
		}
	}

	/**
	 * 公钥验证签名
	 * 
	 * @param data 原数据
	 * @param sign 数据签名
	 * @param certificatePath 证书路径
	 */
	public static boolean verify(byte[] data, byte[] sign, String certificatePath) {
		// 获得证书（通过数字证书）
		X509Certificate x509 = (X509Certificate) getCertificate(certificatePath);
		try {
			// 由证书构建签名
			Signature signature = Signature.getInstance(x509.getSigAlgName(),PROVIDER);
			// 公钥验证签名
			signature.initVerify(x509);
			signature.update(data);
			return signature.verify(sign);
		} catch (Exception e) {
			throw new RuntimeException("公钥验证签名失败",e);
		}
	}

	/**
	 * 公钥验证签名
	 * 
	 * @param data 原数据
	 * @param sign 数据签名
	 * @param certificatePath 证书路径
	 */
	public static boolean verify(String data, String sign, String certificatePath) {
		try {
			return verify(data.getBytes(StandardCharsets.UTF_8),Base64Util.decode(sign), certificatePath);
		} catch (Exception e) {
			throw new RuntimeException("公钥验证签名失败",e);
		}
	}
	
	// ========================= 基础方法 =========================== //
	
	/**
	 * 通过密钥库获得私钥
	 * 
	 * @param alias 证书别名
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	private static PrivateKey getPrivateKeyByKeyStore(String alias, String keyStorePath, String keyStorePwd,
			boolean isBoth) {
		// 获得密钥库
		KeyStore ks = getKeyStore(keyStorePath, keyStorePwd, isBoth);
		// 获得私钥
		try {
			return (PrivateKey) ks.getKey(alias, keyStorePwd.toCharArray());
		} catch (Exception e) {
			throw new RuntimeException("密钥库别名或密码错误",e);
		}
	}

	/**
	 * 通过数字证书获得公钥
	 * 
	 * @param certificatePath 数字证书路径
	 */
	private static PublicKey getPublicKeyByCertificate(String certificatePath) {
		// 获得证书
		Certificate certificate = getCertificate(certificatePath);
		// 获得公钥
		return certificate.getPublicKey();
	}

	/**
	 * 获得密钥库
	 * 
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	private static KeyStore getKeyStore(String keyStorePath, String keyStorePwd, boolean isBoth) {
		// 实例化密钥库
		KeyStore ks;
		try {
			if (isBoth) {
				// 双向认证使用PKCS12实例化密钥库
				ks = KeyStore.getInstance("PKCS12", PROVIDER);
			} else {
				// 单向认证使用默认方式实例化密钥库
				ks = KeyStore.getInstance(KeyStore.getDefaultType(), PROVIDER);
			}
		} catch (Exception e) {
			throw new RuntimeException("密钥库实例化失败",e);
		}
		// 获得密钥库文件
		FileInputStream is = null;
		try {
			is = new FileInputStream(keyStorePath);
			// 加载密钥库
			ks.load(is, keyStorePwd.toCharArray());
			return ks;
		} catch (Exception e) {
			throw new RuntimeException("密钥库路径错误或密码错误",e);
		} finally {
			CloseUtil.close(is);
		}
	}

	/**
	 * 获得数字证书(通过数字证书)
	 * 
	 * @param certificatePath 数字证书路径
	 */
	private static Certificate getCertificate(String certificatePath) {
		try {
			// 实例化证书工厂
			CertificateFactory cf = CertificateFactory.getInstance("X.509",PROVIDER);
			// 取得证书文件流
			FileInputStream is = new FileInputStream(certificatePath);
			// 生成证书
			Certificate certificate = cf.generateCertificate(is);
			is.close();
			return certificate;
		} catch (Exception e) {
			throw new RuntimeException("数据库书路径错误",e);
		}
	}

	/**
	 * 获得数字证书(通过密钥库)
	 * 
	 * @param alias 证书别名
	 * @param keyStorePath 密钥库路径
	 * @param keyStorePwd 密钥库密码
	 * @param isBoth 是否为双向认证
	 */
	private static Certificate getCertificate(String alias, String keyStorePath, String keyStorePwd, boolean isBoth) {
		// 获得密钥库
		KeyStore ks = getKeyStore(keyStorePath, keyStorePwd, isBoth);
		// 获得证书
		try {
			return ks.getCertificate(alias);
		} catch (KeyStoreException e) {
			throw new RuntimeException("证书别名错误",e);
		}
	}
}
